#include <iostream>
#include <string>

#include "mpc_quadrotor_mujoco/QuadrotorInterface.h"
#include "mpc_quadrotor_mujoco/QuadrotorSystemDynamics.h"

#include <ocs2_core/cost/QuadraticStateCost.h>
#include <ocs2_core/cost/QuadraticStateInputCost.h>
#include <ocs2_core/initialization/OperatingPoints.h>
#include <ocs2_core/misc/LoadData.h>

#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/path.hpp>

namespace ocs2 {
    namespace quadrotor {
        
        QuadrotorInterface::QuadrotorInterface(const std::string& taskFile, const std::string& libraryFolder) {
            boost::filesystem::path taskFilePath(taskFile);
            if (boost::filesystem::exists(taskFilePath)) {
                std::cerr << "[QuadrotorInterface] Loading task file: " << taskFilePath << std::endl;
            } else {
                throw std::invalid_argument("[QuadrotorInterface] Task file not found: " + taskFilePath.string());
            }

            boost::filesystem::path libraryFolderPath(libraryFolder);
            boost::filesystem::create_directories(libraryFolderPath);
            std::cerr << "[QuadrotorInterface] Generated library path: " << libraryFolderPath << std::endl;

            loadData::loadEigenMatrix(taskFile, "initialState", initialState_);
            std::cerr << "x_init: " << initialState_.transpose() << std::endl;

            ddpSettings_ = ddp::loadSettings(taskFile, "ddp");
            mpcSettings_ = mpc::loadSettings(taskFile, "mpc");

            referenceManagerPtr_.reset(new ReferenceManager);

            matrix_t Q(STATE_DIM, STATE_DIM);
            matrix_t R(INPUT_DIM, INPUT_DIM);
            matrix_t Qf(STATE_DIM, STATE_DIM);
            loadData::loadEigenMatrix(taskFile, "Q", Q);
            loadData::loadEigenMatrix(taskFile, "R", R);
            loadData::loadEigenMatrix(taskFile, "Q_final", Qf);
            std::cerr << "Q:    \n" << Q << std::endl;
            std::cerr << "R:    \n" << R << std::endl;
            std::cerr << "Q_final:    \n" << Qf << std::endl;

            problem_.costPtr->add("cost", std::make_unique<QuadraticStateInputCost>(Q, R));
            problem_.finalCostPtr->add("finalCost", std::make_unique<QuadraticStateCost>(Qf));

            auto quadrotorParameters = quadrotor::loadSettings(taskFile, "QuadrotorParameters", true);
            problem_.dynamicsPtr.reset(new QuadrotorSystemDynamics(quadrotorParameters));

            auto rolloutSettings = rollout::loadSettings(taskFile, "rollout");
            rolloutPtr_.reset(new TimeTriggeredRollout(*problem_.dynamicsPtr, rolloutSettings));

            vector_t initialInput = vector_t::Zero(INPUT_DIM);
            initialInput(0) = quadrotorParameters.quadrotorMass_ * quadrotorParameters.gravity_;
            operatingPointPtr_.reset(new OperatingPoints(initialState_, initialInput));
        }
    } // namespace quadrotor
} // namespace ocs2

